package sf.database.datasource;

import sf.spring.util.Assert;
import sf.spring.util.CollectionUtils;

import javax.sql.DataSource;
import java.util.AbstractMap;
import java.util.Map;
import java.util.Random;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 简易主从数据源管理,使用spring的 RoutingDataSource 过于复杂.
 */
public class SimpleRoutingDataSource extends ADataSource {
    /**
     * 数据源
     */
    protected final Map.Entry<String, DataSource> master;

    protected final Map<String, DataSource> slaves = new ConcurrentHashMap<>(2);

    /**
     * 构造方法
     * @param masterKey        主数据源key
     * @param masterDataSource 主数据源
     * @param slaves           从数据源
     */
    public SimpleRoutingDataSource(String masterKey, DataSource masterDataSource, Map<String, DataSource> slaves) {
        Assert.notNull(masterKey, "");
        Assert.notNull(masterDataSource, "");
        this.master = new AbstractMap.SimpleEntry<>(masterKey, masterDataSource);
        if (CollectionUtils.isNotEmpty(slaves)) {
            this.slaves.putAll(slaves);
        }
    }

    public DataSource getSlavesDataSources() {
        if (slaves.isEmpty()) {
            return master.getValue();
        }
        return nextSlave();
    }

    protected DataSource nextSlave() {
        //随机，todo，换成顺序
        int random = new Random().nextInt(slaves.size());
        int i = 0;
        for (Map.Entry<String, DataSource> entry : slaves.entrySet()) {
            if (i == random) {
                return entry.getValue();
            }
            i++;
        }
        return null;
    }

    public void addSlavesDataSource(String key, DataSource ds) {
        slaves.put(key, ds);
    }

    public DataSource getDataSource(String key) {
        if (key == null) {
            return master.getValue();
        }
        if (key.equals(master.getKey())) {
            return master.getValue();
        } else {
            return slaves.get(key);
        }
    }

    public DataSource getMasterDataSource() {
        return master.getValue();
    }

    public Map.Entry<String, DataSource> getMaster() {
        return master;
    }

    public Map<String, DataSource> getSlaves() {
        return slaves;
    }

    @Override
    protected DataSource targetDataSource() {
        return getDataSource(null);
    }
}
